utils: handle flash_attn missing from importlib packages_distributions without crashing by SAY-5 · Pull Request #45524 · huggingface/transformers

SAY-5 · 2026-04-20T05:44:53Z

is_flash_attn_2_available, is_flash_attn_3_available, is_flash_attn_4_available, and is_flash_attn_greater_or_equal all do two checks:

is_available, _ = _is_package_available("flash_attn", return_version=True)
is_available = is_available and "flash-attn" in [
    pkg.replace("_", "-") for pkg in PACKAGE_DISTRIBUTION_MAPPING["flash_attn"]
]

Step 1 uses importlib.util.find_spec, which returns a spec if any flash_attn import is findable, an editable install, a namespace package, a bundled shim, or a stub module sitting under another project. Step 2 then assumes that every findable import name also has an entry in importlib.metadata.packages_distributions().

That assumption does not hold in practice. On Python 3.13 with ComfyUI setups (as reported in #45520), and more generally in any environment where the flash_attn import is resolvable via a non-pip source, packages_distributions() has no "flash_attn" key. Because the list comprehension is evaluated before the in operator, short-circuit evaluation of the outer and does not protect us, the KeyError fires during transformers import and takes down the whole process before any model is loaded.

Swap the four raising subscripts for .get(name, []). If the name is missing from the distribution map we simply conclude that the requested flash-attention flavour is not properly installed, which is the answer is_flash_attn_*_available() would have returned anyway, instead of raising. The inner helper _is_package_available already wraps the same subscript in a try/except, so we are only making the outer call sites match that contract.

…not in the distribution map is_flash_attn_2_available / _3 / _4 / _greater_or_equal do two checks: is_available, _ = _is_package_available("flash_attn", return_version=True) is_available = is_available and "flash-attn" in [ pkg.replace("_", "-") for pkg in PACKAGE_DISTRIBUTION_MAPPING["flash_attn"] ] Step 1 uses importlib.util.find_spec, which returns a spec if any "flash_attn" import is findable (an editable install, a namespace package, a bundled shim, or a stub module under another project). Step 2 then assumes that every findable import name also has an entry in importlib.metadata.packages_distributions(). That assumption does not hold. On Python 3.13 with ComfyUI setups (huggingface#45520), and in any environment where the import is resolvable via a non-pip source, packages_distributions() has no "flash_attn" key. Because the list comprehension is evaluated before the `in` operator, short-circuit evaluation of the outer `and` does not protect us - the KeyError fires during `transformers` import and takes down the whole process before any model is loaded. Swap the four raising subscripts for `.get(name, [])`. If the name is missing from the distribution map we simply conclude that the requested flash-attention flavour is not properly installed - which is the same answer is_flash_attn_*_available() would have returned anyway - instead of raising. The inner helper `_is_package_available` already wraps the same subscript in a try/except, so we are only making the outer call sites match that contract. Fixes huggingface#45520

Rocketknight1 · 2026-04-21T12:16:47Z

cc @vasqu

vasqu

Thank you, seems valid in either case. Can you

Update the same in modeling flash attn utils

transformers/src/transformers/modeling_flash_attention_utils.py

Lines 73 to 109 in 57f9936

    
           FLASH_ATTENTION_COMPATIBILITY_MATRIX = { 
        
               2: { 
        
                   "flash_attn_version": 2, 
        
                   "general_availability_check": is_flash_attn_2_available, 
        
                   "pkg_availability_check": lambda *args, **kwargs: importlib.util.find_spec("flash_attn") is not None 
        
                   and "flash-attn" in [pkg.replace("_", "-") for pkg in PACKAGE_DISTRIBUTION_MAPPING["flash_attn"]], 
        
                   "supported_devices": ( 
        
                       (is_torch_cuda_available, "cuda"), 
        
                       (is_torch_mlu_available, "mlu"), 
        
                       (is_torch_npu_available, "npu"), 
        
                       (is_torch_xpu_available, "xpu"), 
        
                   ), 
        
                   "custom_supported_devices": ( 
        
                       (is_torch_npu_available, "Detect using FlashAttention2 on Ascend NPU."), 
        
                       ( 
        
                           is_torch_xpu_available, 
        
                           f"Detect using FlashAttention2 (via kernel `{FLASH_ATTN_KERNEL_FALLBACK['flash_attention_2']}`) on XPU.", 
        
                       ), 
        
                   ), 
        
               }, 
        
               3: { 
        
                   "flash_attn_version": 3, 
        
                   "general_availability_check": is_flash_attn_3_available, 
        
                   "pkg_availability_check": lambda *args, **kwargs: importlib.util.find_spec("flash_attn_interface") is not None 
        
                   and "flash-attn-3" in [pkg.replace("_", "-") for pkg in PACKAGE_DISTRIBUTION_MAPPING["flash_attn_interface"]], 
        
                   "supported_devices": ((is_torch_cuda_available, "cuda"),), 
        
                   "cuda_min_major_version": 8,  # Ampere 
        
               }, 
        
               4: { 
        
                   "flash_attn_version": 4, 
        
                   "general_availability_check": is_flash_attn_4_available, 
        
                   "pkg_availability_check": lambda *args, **kwargs: importlib.util.find_spec("flash_attn") is not None 
        
                   and "flash-attn-4" in [pkg.replace("_", "-") for pkg in PACKAGE_DISTRIBUTION_MAPPING["flash_attn"]], 
        
                   "supported_devices": ((is_torch_cuda_available, "cuda"),), 
        
                   "cuda_min_major_version": 9,  # Hopper 
        
               }, 
        
           }

Add a small test similar to

transformers/tests/utils/test_modeling_utils.py

Line 2840 in 57f9936

def test_not_available_flash(self):

…gression tests Signed-off-by: say <say.apm35@gmail.com>

SAY-5 · 2026-05-25T05:04:06Z

Applied both requests:

Extended the same .get() guard to the three pkg_availability_check lambdas in modeling_flash_attention_utils.py (lines 78, 97, 105 in the matrix).
Added TestFlashAttnDistributionMapMissing in tests/utils/test_import_utils.py with three unit tests that patch PACKAGE_DISTRIBUTION_MAPPING to omit flash_attn/flash_attn_interface and assert the availability functions return False rather than raising KeyError.

Signed-off-by: Sai Asish Y <say.apm35@gmail.com>

vasqu · 2026-05-25T15:37:23Z

    assert modeling_auto.__name__ == "transformers.models.auto.modeling_auto"
+
+
+class TestFlashAttnDistributionMapMissing(unittest.TestCase):


Let's remove these tests please and add a test to modeling utils, right below

transformers/tests/utils/test_modeling_utils.py

Line 2840 in 57f9936

def test_not_available_flash(self):

Where we check that no error is there anymore

…ng_utils Signed-off-by: Sai Asish Y <say.apm35@gmail.com>

SAY-5 · 2026-05-25T22:23:55Z

Moved the distribution-map regression test to TestAttentionImplementation in tests/utils/test_modeling_utils.py (right below test_not_available_flash) and removed it from test_import_utils.py.

SAY-5 · 2026-05-26T01:38:33Z

Both failing jobs are pre-existing flakes unrelated to this change.

tests_processors: 7 tests fail with AttributeError: module 'torchvision.io' has no attribute 'read_video' - torchvision API regression, not in scope of this PR (no processor or torchvision code touched).

tests_tensor_parallel_ci: 2 tests fail with ProcessExitedException: process 0 terminated with signal SIGABRT on cwm and solar_open models - same failure visible on PRs #46142, #45919, #46146 at the same time, all unrelated to flash-attn utils.

vasqu

Should be good now, thanks 🤗 gh is having trouble again not sure if I can merge now or later but ping me if you feel like I forgot this PR

Edit: yea will wait a bit, some stuff is failing because of the gh issues

…s without crashing (huggingface#45524) * utils: stop crashing with KeyError when flash_attn is importable but not in the distribution map is_flash_attn_2_available / _3 / _4 / _greater_or_equal do two checks: is_available, _ = _is_package_available("flash_attn", return_version=True) is_available = is_available and "flash-attn" in [ pkg.replace("_", "-") for pkg in PACKAGE_DISTRIBUTION_MAPPING["flash_attn"] ] Step 1 uses importlib.util.find_spec, which returns a spec if any "flash_attn" import is findable (an editable install, a namespace package, a bundled shim, or a stub module under another project). Step 2 then assumes that every findable import name also has an entry in importlib.metadata.packages_distributions(). That assumption does not hold. On Python 3.13 with ComfyUI setups (huggingface#45520), and in any environment where the import is resolvable via a non-pip source, packages_distributions() has no "flash_attn" key. Because the list comprehension is evaluated before the `in` operator, short-circuit evaluation of the outer `and` does not protect us - the KeyError fires during `transformers` import and takes down the whole process before any model is loaded. Swap the four raising subscripts for `.get(name, [])`. If the name is missing from the distribution map we simply conclude that the requested flash-attention flavour is not properly installed - which is the same answer is_flash_attn_*_available() would have returned anyway - instead of raising. The inner helper `_is_package_available` already wraps the same subscript in a try/except, so we are only making the outer call sites match that contract. Fixes huggingface#45520 * utils: fix same KeyError in modeling_flash_attention_utils and add regression tests Signed-off-by: say <say.apm35@gmail.com> * style: fix ruff formatting in flash_attn fix and tests Signed-off-by: Sai Asish Y <say.apm35@gmail.com> * test: move flash_attn distribution-map regression test to test_modeling_utils Signed-off-by: Sai Asish Y <say.apm35@gmail.com> * revert this * is gh alright? --------- Signed-off-by: say <say.apm35@gmail.com> Signed-off-by: Sai Asish Y <say.apm35@gmail.com> Co-authored-by: SAY-5 <SAY-5@users.noreply.github.com> Co-authored-by: vasqu <antonprogamer@gmail.com> Co-authored-by: Anton Vlasjuk <73884904+vasqu@users.noreply.github.com>

vasqu reviewed Apr 23, 2026

View reviewed changes

Rocketknight1 mentioned this pull request Apr 27, 2026

Fix KeyError for flash_attn in import_utils.py on Python 3.13Fix #45650

Closed

3 tasks

utils: fix same KeyError in modeling_flash_attention_utils and add re…

6e7cd69

…gression tests Signed-off-by: say <say.apm35@gmail.com>

style: fix ruff formatting in flash_attn fix and tests

3f8ef40

Signed-off-by: Sai Asish Y <say.apm35@gmail.com>

vasqu reviewed May 25, 2026

View reviewed changes

test: move flash_attn distribution-map regression test to test_modeli…

282602d

…ng_utils Signed-off-by: Sai Asish Y <say.apm35@gmail.com>

vasqu added 3 commits May 26, 2026 14:24

Merge branch 'main' into fix/flash-attn-keyerror-45520

2a8593a

revert this

b66b877

is gh alright?

e931cfd

vasqu approved these changes May 26, 2026

View reviewed changes

Merge branch 'main' into fix/flash-attn-keyerror-45520

e19a994

vasqu enabled auto-merge May 26, 2026 14:40

vasqu added this pull request to the merge queue May 26, 2026

github-merge-queue Bot removed this pull request from the merge queue due to failed status checks May 26, 2026

vasqu added this pull request to the merge queue May 26, 2026

Merged via the queue into huggingface:main with commit 02af5bf May 26, 2026
30 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

utils: handle flash_attn missing from importlib packages_distributions without crashing - #45524

utils: handle flash_attn missing from importlib packages_distributions without crashing#45524
vasqu merged 8 commits into
huggingface:mainfrom
SAY-5:fix/flash-attn-keyerror-45520

SAY-5 commented Apr 20, 2026 •

edited

Loading

Uh oh!

Rocketknight1 commented Apr 21, 2026

Uh oh!

vasqu left a comment

Uh oh!

SAY-5 commented May 25, 2026

Uh oh!

vasqu May 25, 2026

Uh oh!

SAY-5 commented May 25, 2026

Uh oh!

SAY-5 commented May 26, 2026

Uh oh!

vasqu left a comment •

edited

Loading

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

	FLASH_ATTENTION_COMPATIBILITY_MATRIX = {
	2: {
	"flash_attn_version": 2,
	"general_availability_check": is_flash_attn_2_available,
	"pkg_availability_check": lambda args, *kwargs: importlib.util.find_spec("flash_attn") is not None
	and "flash-attn" in [pkg.replace("_", "-") for pkg in PACKAGE_DISTRIBUTION_MAPPING["flash_attn"]],
	"supported_devices": (
	(is_torch_cuda_available, "cuda"),
	(is_torch_mlu_available, "mlu"),
	(is_torch_npu_available, "npu"),
	(is_torch_xpu_available, "xpu"),
	),
	"custom_supported_devices": (
	(is_torch_npu_available, "Detect using FlashAttention2 on Ascend NPU."),
	(
	is_torch_xpu_available,
	f"Detect using FlashAttention2 (via kernel `{FLASH_ATTN_KERNEL_FALLBACK['flash_attention_2']}`) on XPU.",
	),
	),
	},
	3: {
	"flash_attn_version": 3,
	"general_availability_check": is_flash_attn_3_available,
	"pkg_availability_check": lambda args, *kwargs: importlib.util.find_spec("flash_attn_interface") is not None
	and "flash-attn-3" in [pkg.replace("_", "-") for pkg in PACKAGE_DISTRIBUTION_MAPPING["flash_attn_interface"]],
	"supported_devices": ((is_torch_cuda_available, "cuda"),),
	"cuda_min_major_version": 8, # Ampere
	},
	4: {
	"flash_attn_version": 4,
	"general_availability_check": is_flash_attn_4_available,
	"pkg_availability_check": lambda args, *kwargs: importlib.util.find_spec("flash_attn") is not None
	and "flash-attn-4" in [pkg.replace("_", "-") for pkg in PACKAGE_DISTRIBUTION_MAPPING["flash_attn"]],
	"supported_devices": ((is_torch_cuda_available, "cuda"),),
	"cuda_min_major_version": 9, # Hopper
	},
	}

		assert modeling_auto.__name__ == "transformers.models.auto.modeling_auto"


		class TestFlashAttnDistributionMapMissing(unittest.TestCase):

Uh oh!

Conversation

SAY-5 commented Apr 20, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Rocketknight1 commented Apr 21, 2026

Uh oh!

vasqu left a comment

Choose a reason for hiding this comment

Uh oh!

SAY-5 commented May 25, 2026

Uh oh!

vasqu May 25, 2026

Choose a reason for hiding this comment

Uh oh!

SAY-5 commented May 25, 2026

Uh oh!

SAY-5 commented May 26, 2026

Uh oh!

vasqu left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

SAY-5 commented Apr 20, 2026 •

edited

Loading

vasqu left a comment •

edited

Loading